from typing import Optional, List

from peewee import Model, CharField, BooleanField, IntegerField, FloatField
from playhouse.shortcuts import model_to_dict
from pydantic import BaseModel

from db.my_sqlite import DB


# 定义文件模型
class Web(Model):
    base_url = CharField(help_text="站点域名")
    userid = IntegerField(help_text="用户id")
    rag_id = CharField(help_text="rag的id")
    is_chunk = BooleanField(help_text="是否分块")
    sub_urls = CharField(help_text="域名下网页的url，list转str", default="")
    dis_suffix = CharField(help_text="过滤后缀", default="")
    dis_url = CharField(help_text="过滤网址", default="")
    progress = FloatField(help_text="分块进度", default=0)
    create_date = CharField(help_text="上传日期")
    stop = BooleanField(help_text="是否停用", default=False)
    id = IntegerField(primary_key=True, help_text="主键id")
    # 0 文件, 1 评测集
    type = IntegerField(help_text="文件类型", default=0)

    class Meta:
        database = DB


class WebModel(BaseModel):
    base_url: Optional[str]
    userid: Optional[int]
    rag_id: Optional[str]
    is_chunk: Optional[bool]
    sub_urls: Optional[str]
    dis_suffix: Optional[str]
    dis_url: Optional[str]
    progress: Optional[float]
    create_date: Optional[str]
    stop: Optional[bool]
    id: Optional[int]
    type: Optional[int]

class WebsTable:
    def __init__(self, db):
        self.db = db
        self.db.create_tables([Web])

    @staticmethod
    def delete_base_url(userid: int, base_url: str, rag_id: str) -> bool:
        return Web.delete().where(
            (Web.userid == userid) &
            (Web.base_url == base_url) &
            (Web.rag_id == rag_id)
        ).execute() > 0

    @staticmethod
    def update_single(userid: int, base_url: str, rag_id: str, single_url: str) -> bool:
        return Web.update(sub_urls=single_url).where(
            (Web.userid == userid) &
            (Web.base_url == base_url) &
            (Web.rag_id == rag_id)
        ).execute() > 0

    @staticmethod
    def get_web(userid: int, base_url: str, rag_id: str) -> Optional[WebModel]:
        file = Web.get_or_none(
            (Web.userid == userid) &
            (Web.base_url == base_url) &
            (Web.rag_id == rag_id)
        )
        if file:
            return WebModel(**model_to_dict(file))
        return None

    @staticmethod
    def get_webs_by_rag_id(userid: int, rag_id: str,query: str) -> List[WebModel]:
        base_query = Web.select().where(
            (Web.userid == userid) &
            (Web.rag_id == rag_id))
        if query:
            base_query = base_query.where(Web.base_url ** f"%{query}%")  # 添加模糊查询条件

        return [WebModel(**model_to_dict(file)) for file in base_query]

    @staticmethod
    def get_webs_count_by_rag_id(userid: int, rag_id: str,query: str) -> int:
        base_query = (Web.select().where(
            (Web.userid == userid) &
            (Web.rag_id == rag_id)))
        if query:
            base_query = base_query.where(Web.base_url ** f"%{query}%")  # 添加模糊查询条件
        count = base_query.count()
        return count

    @staticmethod
    def insert_new_web(
            userid: int,
            base_url: str,
            rag_id: str,
            dis_url: str,
            dis_suffix: str,
            create_date: str,
            is_chunk: bool = False,
            progress: float = 0,
            type: int = 1,
    ) -> Optional[WebModel]:
        file = Web.create(
            userid=userid,
            base_url=base_url,
            rag_id=rag_id,
            dis_url=dis_url,
            dis_suffix=dis_suffix,
            is_chunk=is_chunk,
            progress=progress,
            create_date=create_date,
            stop=False,
            type=type,
        )
        return file

    @staticmethod
    def update_web(
            userid: int,
            base_url: str,
            rag_id: str,
            dis_url: str,
            dis_suffix: str,
            is_chunk: bool = False,
            progress: float = 0
    ) -> Optional[bool]:
        res = Web.update(
            dis_url=dis_url,
            dis_suffix=dis_suffix,
            is_chunk=is_chunk,
            sub_urls="",
            progress=progress
        ).where(
            (Web.userid == userid) &
            (Web.base_url == base_url) &
            (Web.rag_id == rag_id)
        ).execute() > 0

        return res

    @staticmethod
    def update_chunk_status(
            userid: int,
            base_url: str,
            rag_id: str,
            is_chunk: bool = False
    ) -> Optional[bool]:
        res = Web.update(
            is_chunk=is_chunk,
        ).where(
            (Web.userid == userid) &
            (Web.base_url == base_url) &
            (Web.rag_id == rag_id)
        ).execute()>0

        return res

    @staticmethod
    async def update_webs_chunk_progress(
            userid: int,
            base_url: str,
            rag_id: str,
            progress: float,
    ) -> bool:
        res = Web.update(
            progress=progress
        ).where(
            (Web.userid == userid) &
            (Web.base_url == base_url) &
            (Web.rag_id == rag_id)
        ).execute() > 0

        return res

    @staticmethod
    async def get_webs_chunk_progress(
            userid: int,
            base_url: str,
            rag_id: str,
    ) -> Optional[float]:
        return Web.get_or_none(
            (Web.userid == userid) &
            (Web.base_url == base_url) &
            (Web.rag_id == rag_id)
        ).progress
    @staticmethod
    def update_web_status(
            web_id: int,
            stop: bool ,
    ) -> bool:
        res = Web.update(
            stop=stop,
        ).where(
            (Web.id == web_id)
        ).execute() > 0
        return res

Webs = WebsTable(DB)
